#include <iostream>
#include <vector>
#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/kdtree/kdtree_flann.h>
#include <Eigen/Core>
#include <Eigen/Dense>
#include <Utils.h>
using namespace std;
using namespace Eigen;
namespace xio=lyxutils::io;
namespace xos=lyxutils::os;
namespace np=lyxutils::numpy;

void print_usage(){
    cout<<"usage:transPC pc.txt T.txt [ref_pts.txt align_pts.txt] [-o result.txt]"<<endl;
    cout<<"explain:transform source point cloud from system src to system dest"<<endl;
    cout<<"            and apply distortion if last two parameters are offered"<<endl;
    cout<<"parameters:"<<endl;
    cout<<"  pc.txt         =>the source point cloud"<<endl;
    cout<<"  T.txt          =>the transformation matrix (4x4) from dest to src(if not reversed)"<<endl;
    cout<<"                   namely T_src_dest,usually derived from `estimateT src.txt dest.txt`"<<endl;
    cout<<"  ref_pts.txt    =>reference point set in source point cloud(for distortion adjust)"<<endl;
    cout<<"  align_pts.txt  =>align point set in destination point cloud(for distortion adjust)"<<endl;
    cout<<"arguments:"<<endl;
    cout<<"  -h    print this help message"<<endl;
    cout<<"  -r    take the input transformation matrix as reversed (namely from system src to dest=>T_{dest,src})"<<endl;
    cout<<"  -m    maximum neighboring anchor point to use(default is 4)"<<endl;
    cout<<"  -o    set the transformed [and distorted] point cloud output file name"<<endl;
    cout<<"  -f    force overwrite if the output file already exist!"<<endl;
    cout<<"  -d    use distance instead of squared distance to calculate weights"<<endl;
}

int main(int argc,char **argv) {
    xio::CLParser parser;
    try{
        parser.parse(argc,argv);
    }catch(invalid_argument ia){
        cout<<"invalid argument:"<<ia.what()<<endl;
        print_usage();
        return 1;
    }
    //io operations
    vector<vector<float> > pc_raw;
    vector<vector<float> > T_raw;//T_12
    vector<vector<float> > ref_raw;
    vector<vector<float> > align_raw;
    vector<string> paras=parser.getParameters();
    if((paras.size()!=2&&paras.size()!=4)||parser.hasOption("h")){
        print_usage();
        return 1;
    }
    bool disAdjust=(paras.size()==4);//if last two parameters offered,distortion adjust is enabled
    for(int i=0;i<paras.size();++i){
        if(!xos::file(xos::abspath(paras[i]))){
            cout<<"file:"<<paras[i]<<" not exist!"<<endl;
            return 1;
        }
    }
    int max_neighbor=parser.hasOption("m")?stoi(parser.getOptionArg("m")):4;
    if(max_neighbor<0){
        cout<<"neighbors cannot be less than 0!"<<endl;
        return 1;
    }
    string outputFile=parser.hasOption("o")?parser.getOptionArg("o"):"result.txt";
    if(xos::file(xos::abspath(outputFile))&&!parser.hasOption("f")){
        cout<<"the outputFile name:"<<outputFile<<" already exist,using -f to force overwrite"<<endl;
        return 1;
    }
    xio::read_csv(paras[0], pc_raw);//read original pc_raw
    xio::read_csv(paras[1],T_raw);
    if(disAdjust){
        xio::read_csv(paras[2],ref_raw);
        xio::read_csv(paras[3],align_raw);
        if(max_neighbor>ref_raw.size()){
            max_neighbor=ref_raw.size();
        }
    }
    auto test_pc=[&](const vector<vector<float> > &pc)->bool{
        if(pc.size()==0||pc[0].size()<3){
            cout<<"point cloud file should has raws>0 and cols>=3,current (row,col)=(";
            cout<<pc.size()<<",";
            if(pc.size()>0)cout<<pc[0].size();
            cout<<")"<<endl;
            return false;
        }
        return true;
    };
    if(!test_pc(pc_raw))return 1;
    if(disAdjust){
        if(!test_pc(ref_raw))return 1;
        if(!test_pc(align_raw))return 1;
    }
    if(T_raw.size()!=4||T_raw[0].size()!=4)cout<<"transformation matrix should be 4x4!"<<endl;
    //real work here
    //inverse T if needed
    if(!parser.hasOption("r")){//if not reversed, need to get the inverse of T12
        Matrix4d t;
        for(int i=0;i<4;++i)for(int j=0;j<4;++j)t(i,j)=T_raw[i][j];
        Matrix4d t1=t.inverse();
        for(int i=0;i<4;++i)for(int j=0;j<4;++j)T_raw[i][j]=t1(i,j);
    }
    //transform
    vector<vector<float> > pc_trans(pc_raw.size(),vector<float>(pc_raw[0].size(),0));
    for(int i=0;i<pc_trans.size();++i){//index of point
        for(int j=0;j<3;++j){//index of coord
            pc_trans[i][j]=T_raw[j][0]*pc_raw[i][0]+T_raw[j][1]*pc_raw[i][1]+T_raw[j][2]*pc_raw[i][2]+T_raw[j][3];
        }
        for(int j=3;j<pc_trans[i].size();++j){
            pc_trans[i][j]=pc_raw[i][j];
        }
    }
    // xio::write_csv("pc_to_photo.txt",pc_trans,",");//test output
    //estimate distortion parameters
    if(!disAdjust)max_neighbor=0;
    //distort
    if(max_neighbor!=0){//if max number of neighboring anchor point is 0,no need for distort
        vector<vector<float> > ref_trans(ref_raw.size(),vector<float>(ref_raw[0].size(),0));
        vector<vector<float> > diff(ref_raw.size(),vector<float>(ref_raw[0].size(),0));
        for(int i=0;i<ref_trans.size();++i){
            for(int j=0;j<ref_trans[i].size();++j){
                ref_trans[i][j]=T_raw[j][0]*ref_raw[i][0]+T_raw[j][1]*ref_raw[i][1]+T_raw[j][2]*ref_raw[i][2]+T_raw[j][3];
                diff[i][j]=align_raw[i][j]-ref_trans[i][j];
            }
        }
        pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_ref(new pcl::PointCloud<pcl::PointXYZ>);
        pcl::KdTreeFLANN<pcl::PointXYZ>::Ptr tree(new pcl::KdTreeFLANN<pcl::PointXYZ>);
        for(int i=0;i<ref_trans.size();++i){
            cloud_ref->points.push_back(pcl::PointXYZ(ref_trans[i][0],ref_trans[i][1],ref_trans[i][2]));
        }
        cloud_ref->width=ref_trans.size();
        cloud_ref->height=1;
        cloud_ref->is_dense=true;
        tree->setInputCloud(cloud_ref);
        vector<vector<float> > distort_vec(pc_trans.size(),vector<float>(3,0));
        for(int i=0;i<pc_trans.size();++i){
            vector<int> idx_neighbors(max_neighbor);
            vector<float> sqr_dis(max_neighbor);
            int num_neighbors=0;
            if(max_neighbor==ref_trans.size()){//no need to search for neighbors,using all ref points
                num_neighbors=max_neighbor;
                np::range(idx_neighbors,max_neighbor);
                for(int j=0;j<max_neighbor;++j){
                    float dx=pc_trans[i][0]-ref_trans[j][0];
                    float dy=pc_trans[i][1]-ref_trans[j][1];
                    float dz=pc_trans[i][2]-ref_trans[j][2];
                    sqr_dis[j]=dx*dx+dy*dy+dz*dz;
                    idx_neighbors[j]=j;//use every anchor point as neighbor
                }
            }
            else{
                num_neighbors=tree->nearestKSearch(pcl::PointXYZ(pc_trans[i][0],pc_trans[i][1],pc_trans[i][2]),max_neighbor,idx_neighbors,sqr_dis);
            }
            if(num_neighbors>0){
                if(parser.hasOption("d")){//use distance instead of squared distance
                    np::apply(sqr_dis,[&](float &dis){dis=sqrt(dis);});//get all distances
                }
                vector<double> weights(max_neighbor,1);
                double w_sum=0;
                for(int j=0;j<max_neighbor;++j){
                    for(int k=0;k<max_neighbor;++k){
                        if(k!=j){
                            weights[j]*=sqr_dis[k];
                        }
                    }
                    w_sum+=weights[j];
                }
                np::apply(weights,[&w_sum](double &w){w/=w_sum;});
                for(int j=0;j<max_neighbor;++j){
                    for(int k=0;k<3;++k)distort_vec[i][k]+=diff[idx_neighbors[j]][k]*weights[j];
                }
            }
            for(int j=0;j<3;++j)pc_trans[i][j]+=distort_vec[i][j];
        }
    }
    //save to file
    xio::write_csv(outputFile,pc_trans,",");
    return 0;
}
